import tkinter as tk
from tkinter import ttk, scrolledtext, messagebox
import subprocess
import threading
import time
import re
import sys
import os
import zipfile
import io
from datetime import datetime

class Iperf3TesterApp:
    def __init__(self, root):
        self.root = root
        self.root.title("Iperf3网络测试工具")
        self.root.geometry("900x700")
        self.root.minsize(800, 600)
        
        # 设置中文字体支持
        self.setup_fonts()
        
        # 测试进程和线程控制
        self.test_process = None
        self.is_testing = False
        self.test_thread = None
        
        # iperf3相关设置
        self.detected_iperf_command = None
        self.use_builtin_iperf3 = True  # 默认使用内置版本
        self.builtin_iperf3_path = None
        
        # 创建主布局
        self.create_widgets()
    
    def setup_fonts(self):
        # 在Windows上设置合适的字体
        self.default_font = ("SimHei", 10)
        self.header_font = ("SimHei", 12, "bold")
        self.monospace_font = ("Courier New", 10)
        self.speed_font = ("SimHei", 36, "bold")  # 大字体显示速度
    
    def create_widgets(self):
        # 创建主框架
        main_frame = ttk.Frame(self.root, padding="10")
        main_frame.pack(fill=tk.BOTH, expand=True)
        
        # 模式选择区域
        mode_frame = ttk.LabelFrame(main_frame, text="测试模式", padding="8")
        mode_frame.pack(fill=tk.X, pady=(0, 10))
        
        self.mode_var = tk.StringVar(value="server")
        server_radio = ttk.Radiobutton(mode_frame, text="服务端模式", variable=self.mode_var, 
                                     value="server", command=self.on_mode_change)
        server_radio.pack(side=tk.LEFT, padx=20, pady=5)
        
        client_radio = ttk.Radiobutton(mode_frame, text="客户端模式", variable=self.mode_var, 
                                     value="client", command=self.on_mode_change)
        client_radio.pack(side=tk.LEFT, padx=20, pady=5)
        
        # 客户端设置区域
        self.client_frame = ttk.LabelFrame(main_frame, text="客户端设置", padding="8")
        self.client_frame.pack(fill=tk.X, pady=(0, 10))
        
        # 服务器IP
        ip_frame = ttk.Frame(self.client_frame)
        ip_frame.pack(fill=tk.X, pady=(0, 5))
        
        ip_label = ttk.Label(ip_frame, text="服务器IP:")
        ip_label.pack(side=tk.LEFT, padx=(0, 10))
        
        self.server_ip_entry = ttk.Entry(ip_frame, width=20)
        self.server_ip_entry.pack(side=tk.LEFT, padx=(0, 20))
        self.server_ip_entry.insert(0, "127.0.0.1")
        
        # 端口
        port_label = ttk.Label(ip_frame, text="端口:")
        port_label.pack(side=tk.LEFT, padx=(0, 10))
        
        self.port_entry = ttk.Entry(ip_frame, width=10)
        self.port_entry.pack(side=tk.LEFT, padx=(0, 20))
        self.port_entry.insert(0, "5201")
        
        # 测试时间
        time_label = ttk.Label(ip_frame, text="测试时间(秒):")
        time_label.pack(side=tk.LEFT, padx=(0, 10))
        
        self.time_entry = ttk.Entry(ip_frame, width=10)
        self.time_entry.pack(side=tk.LEFT)
        self.time_entry.insert(0, "10")
        
        # 高级选项区域
        advanced_frame = ttk.LabelFrame(main_frame, text="高级选项", padding="8")
        advanced_frame.pack(fill=tk.X, pady=(0, 10))
        
        # iperf3来源选择
        iperf_frame = ttk.Frame(advanced_frame)
        iperf_frame.pack(fill=tk.X, pady=(0, 5))
        
        iperf_label = ttk.Label(iperf_frame, text="Iperf3来源:")
        iperf_label.pack(side=tk.LEFT, padx=(0, 10))
        
        self.iperf_source_var = tk.StringVar(value="builtin")
        builtin_radio = ttk.Radiobutton(iperf_frame, text="内置版本", variable=self.iperf_source_var, value="builtin")
        builtin_radio.pack(side=tk.LEFT, padx=10)
        
        system_radio = ttk.Radiobutton(iperf_frame, text="系统安装", variable=self.iperf_source_var, value="system")
        system_radio.pack(side=tk.LEFT, padx=10)
        
        # 协议选择
        proto_frame = ttk.Frame(advanced_frame)
        proto_frame.pack(fill=tk.X, pady=(0, 5))
        
        proto_label = ttk.Label(proto_frame, text="协议:")
        proto_label.pack(side=tk.LEFT, padx=(0, 10))
        
        self.proto_var = tk.StringVar(value="tcp")
        tcp_radio = ttk.Radiobutton(proto_frame, text="TCP", variable=self.proto_var, value="tcp")
        tcp_radio.pack(side=tk.LEFT, padx=10)
        
        udp_radio = ttk.Radiobutton(proto_frame, text="UDP", variable=self.proto_var, value="udp")
        udp_radio.pack(side=tk.LEFT, padx=10)
        
        # 带宽限制 (UDP模式)
        self.bw_frame = ttk.Frame(advanced_frame)
        self.bw_frame.pack(fill=tk.X, pady=(0, 5))
        
        bw_label = ttk.Label(self.bw_frame, text="带宽限制(Mbps):")
        bw_label.pack(side=tk.LEFT, padx=(0, 10))
        
        self.bw_entry = ttk.Entry(self.bw_frame, width=10)
        self.bw_entry.pack(side=tk.LEFT)
        self.bw_entry.insert(0, "100")
        
        # 并行流数
        parallel_frame = ttk.Frame(advanced_frame)
        parallel_frame.pack(fill=tk.X, pady=(0, 5))
        
        parallel_label = ttk.Label(parallel_frame, text="并行流数:")
        parallel_label.pack(side=tk.LEFT, padx=(0, 10))
        
        self.parallel_entry = ttk.Entry(parallel_frame, width=10)
        self.parallel_entry.pack(side=tk.LEFT)
        self.parallel_entry.insert(0, "1")
        
        # 速度显示区域
        speed_frame = ttk.LabelFrame(main_frame, text="当前速度", padding="8")
        speed_frame.pack(fill=tk.X, pady=(0, 10))
        
        self.speed_var = tk.StringVar(value="0.00 Mbits/sec")
        speed_label = ttk.Label(speed_frame, textvariable=self.speed_var, font=self.speed_font)
        speed_label.pack(pady=20)
        
        # 操作按钮区域
        button_frame = ttk.Frame(main_frame)
        button_frame.pack(fill=tk.X, pady=(0, 10))
        
        self.start_button = ttk.Button(button_frame, text="开始测试", command=self.start_test)
        self.start_button.pack(side=tk.LEFT, padx=10)
        
        self.stop_button = ttk.Button(button_frame, text="停止测试", command=self.stop_test, state=tk.DISABLED)
        self.stop_button.pack(side=tk.LEFT, padx=10)
        
        self.clear_button = ttk.Button(button_frame, text="清空日志", command=self.clear_log)
        self.clear_button.pack(side=tk.RIGHT, padx=10)
        
        # 日志输出区域
        log_frame = ttk.LabelFrame(main_frame, text="测试日志", padding="8")
        log_frame.pack(fill=tk.BOTH, expand=True)
        
        self.log_text = scrolledtext.ScrolledText(log_frame, font=self.monospace_font)
        self.log_text.pack(fill=tk.BOTH, expand=True)
        self.log_text.config(state=tk.DISABLED)
        
        # 初始设置
        self.on_mode_change()
    
    def on_mode_change(self):
        # 根据模式启用/禁用相关控件
        mode = self.mode_var.get()
        if mode == "server":
            self._set_frame_state(self.client_frame, tk.DISABLED)
        else:
            self._set_frame_state(self.client_frame, tk.NORMAL)
        
        # 根据协议启用/禁用带宽限制
        proto = self.proto_var.get()
        if proto == "udp":
            self._set_frame_state(self.bw_frame, tk.NORMAL)
        else:
            self._set_frame_state(self.bw_frame, tk.DISABLED)
    
    def _set_frame_state(self, frame, state):
        # 只对支持state属性的组件设置状态
        for child in frame.winfo_children():
            try:
                child.config(state=state)
            except tk.TclError:
                # 如果组件不支持state属性，尝试递归处理其子组件
                if hasattr(child, 'winfo_children'):
                    self._set_frame_state(child, state)
    
    def start_test(self):
        # 清空之前的日志，便于查看新的检测信息
        self.clear_log()
        
        # 检查iperf3是否安装
        iperf3_installed = self.check_iperf3_installed()
        
        if not iperf3_installed:
            # 即使检测失败，也允许用户尝试继续
            response = messagebox.askyesno(
                "警告", 
                "未找到iperf3或无法正确识别。\n\n"
                "请确保：\n"
                "1. iperf3已正确安装\n"
                "2. iperf3已添加到系统PATH环境变量\n"
                "3. 或iperf3.exe位于程序同一目录\n\n"
                "是否仍要尝试运行测试？"
            )
            if not response:
                return
        
        # 禁用开始按钮，启用停止按钮
        self.start_button.config(state=tk.DISABLED)
        self.stop_button.config(state=tk.NORMAL)
        
        # 在单独的线程中运行测试，避免UI冻结
        self.is_testing = True
        self.test_thread = threading.Thread(target=self._run_test)
        self.test_thread.daemon = True
        self.test_thread.start()
    
    def stop_test(self):
        # 停止测试
        self.is_testing = False
        if self.test_process:
            try:
                # 尝试优雅地终止进程
                if sys.platform == 'win32':
                    self.test_process.terminate()
                else:
                    import signal
                    self.test_process.send_signal(signal.SIGINT)
                
                # 等待进程终止
                self.test_process.wait(timeout=2)
            except:
                # 强制终止
                if self.test_process.poll() is None:
                    if sys.platform == 'win32':
                        self.test_process.kill()
                    else:
                        import signal
                        self.test_process.send_signal(signal.SIGKILL)
        
        # 恢复按钮状态
        self.root.after(0, self._update_ui_after_test)
    
    def _update_ui_after_test(self):
        # 更新UI状态
        self.start_button.config(state=tk.NORMAL)
        self.stop_button.config(state=tk.DISABLED)
        self.log("测试已停止")
    
    def _run_test(self):
        try:
            # 确保使用检测到的命令
            if hasattr(self, 'detected_iperf_command') and self.detected_iperf_command:
                iperf_cmd = self.detected_iperf_command
            else:
                # 如果没有检测到，尝试常见名称
                import os
                current_dir = os.path.dirname(os.path.abspath(__file__))
                if os.path.exists(os.path.join(current_dir, "iperf3.exe")):
                    iperf_cmd = os.path.join(current_dir, "iperf3.exe")
                elif os.path.exists(os.path.join(current_dir, "iperf.exe")):
                    iperf_cmd = os.path.join(current_dir, "iperf.exe")
                else:
                    iperf_cmd = "iperf3"  # 默认尝试
            
            self.log(f"使用命令: {iperf_cmd}")
            
            # 构建命令参数
            args = []
            mode = self.mode_var.get()
            
            if mode == "server":
                args.append("-s")
                port = self.port_entry.get().strip()
                if port:
                    args.extend(["-p", port])
                self.log(f"服务端模式已启动，监听端口: {port}")
            else:
                args.append("-c")
                server_ip = self.server_ip_entry.get().strip()
                if not server_ip:
                    self.root.after(0, messagebox.showerror, "错误", "请输入服务器IP")
                    return
                args.append(server_ip)
                
                # 添加端口
                port = self.port_entry.get().strip()
                if port:
                    args.extend(["-p", port])
                
                # 添加测试时间
                test_time = self.time_entry.get().strip()
                if test_time:
                    args.extend(["-t", test_time])
                
                # 添加协议类型
                proto = self.proto_var.get()
                args.extend(["-" + proto])
                
                # 如果是UDP，添加带宽限制
                if proto == "udp":
                    bw = self.bw_entry.get().strip()
                    if bw:
                        args.extend(["-b", bw + "M"])
                
                # 添加并行流数
                parallel = self.parallel_entry.get().strip()
                if parallel and parallel != "1":
                    args.extend(["-P", parallel])
                
                self.log(f"客户端模式已启动，连接到: {server_ip}:{port}")
            
            # 添加实时输出参数
            args.append("-i")
            args.append("1")
            
            # 构建完整命令
            if ' ' in iperf_cmd or os.path.exists(iperf_cmd):
                # 如果是带路径的命令或包含空格，使用shell=True
                full_cmd = f'"{iperf_cmd}" ' + ' '.join(args)
                self.log(f"执行完整命令: {full_cmd}")
                self.test_process = subprocess.Popen(
                    full_cmd,
                    stdout=subprocess.PIPE,
                    stderr=subprocess.STDOUT,
                    text=True,
                    bufsize=1,
                    shell=True
                )
            else:
                # 否则使用列表形式
                full_cmd = [iperf_cmd] + args
                self.log(f"执行命令列表: {full_cmd}")
                self.test_process = subprocess.Popen(
                    full_cmd,
                    stdout=subprocess.PIPE,
                    stderr=subprocess.STDOUT,
                    text=True,
                    bufsize=1
                )
            
            # 实时读取输出
            for line in self.test_process.stdout:
                if not self.is_testing:
                    break
                
                line = line.strip()
                if line:
                    self.log(line)
                    self._parse_speed(line)
            
            # 等待进程结束
            self.test_process.wait()
            
        except Exception as e:
            self.log(f"测试出错: {str(e)}")
        finally:
            self.test_process = None
            self.is_testing = False
            self.root.after(0, self._update_ui_after_test)
    
    def _parse_speed(self, line):
        # 解析速度信息
        # 匹配形如: [  5]   4.00-5.00   sec  1.25 MBytes  10.5 Mbits/sec
        if self.mode_var.get() == "client":
            # 客户端模式解析
            match = re.search(r'(\d+\.\d+)\s+Mbits/sec', line)
            if match:
                speed = match.group(1)
                self.root.after(0, lambda: self.speed_var.set(f"{speed} Mbits/sec"))
        else:
            # 服务端模式解析
            match = re.search(r'(\d+\.\d+)\s+Mbits/sec', line)
            if match:
                speed = match.group(1)
                self.root.after(0, lambda: self.speed_var.set(f"{speed} Mbits/sec"))
    
    def log(self, message):
        # 添加日志信息
        timestamp = datetime.now().strftime("%H:%M:%S")
        log_message = f"[{timestamp}] {message}"
        
        self.root.after(0, self._update_log, log_message)
    
    def _update_log(self, message):
        # 更新日志文本框
        self.log_text.config(state=tk.NORMAL)
        self.log_text.insert(tk.END, message + "\n")
        self.log_text.see(tk.END)
        self.log_text.config(state=tk.DISABLED)
    
    def clear_log(self):
        # 清空日志
        self.log_text.config(state=tk.NORMAL)
        self.log_text.delete("1.0", tk.END)
        self.log_text.config(state=tk.DISABLED)
        self.speed_var.set("0.00 Mbits/sec")
    
    def extract_builtin_iperf3(self):
        """处理内置iperf3可执行文件"""
        try:
            # 由于我们无法直接在代码中嵌入大型二进制文件，我们将采用混合方案：
            # 1. 首先检查程序目录下是否有自带的iperf3.exe
            # 2. 如果没有，检查临时目录
            # 3. 提供明确的下载指导
            
            import tempfile
            temp_dir = tempfile.gettempdir()
            current_dir = os.path.dirname(os.path.abspath(__file__))
            
            # 可能的iperf3路径列表
            possible_paths = [
                os.path.join(current_dir, "iperf3.exe"),  # 程序目录
                os.path.join(current_dir, "tools", "iperf3.exe"),  # tools子目录
                os.path.join(temp_dir, "iperf3.exe")  # 临时目录
            ]
            
            # 检查是否已存在可用的iperf3
            for path in possible_paths:
                if os.path.exists(path) and os.access(path, os.X_OK):
                    self.builtin_iperf3_path = path
                    self.log(f"找到可用的iperf3: {self.builtin_iperf3_path}")
                    return True
            
            # 如果没有找到，提示用户下载
            self.log("提示: 未找到内置的iperf3.exe文件")
            self.log("请按照以下步骤操作:")
            self.log("1. 下载iperf3: https://iperf.fr/iperf-download.php#windows")
            self.log("2. 解压后将iperf3.exe复制到以下位置之一:")
            for path in possible_paths:
                self.log(f"   - {path}")
            
            # 为了方便用户，创建一个tools目录（如果不存在）
            tools_dir = os.path.join(current_dir, "tools")
            if not os.path.exists(tools_dir):
                try:
                    os.makedirs(tools_dir)
                    self.log(f"已创建tools目录: {tools_dir}")
                except:
                    self.log("无法创建tools目录，请手动创建")
            
            return False
            
        except Exception as e:
            self.log(f"处理内置iperf3时出错: {str(e)}")
            return False
    
    def check_iperf3_installed(self):
        # 增强的iperf3检测方法，添加详细日志和更多检测策略，支持内置版本
        self.log("开始检测iperf3...")
        self.detected_iperf_command = None  # 存储找到的有效命令
        
        # 检测策略列表，包含各种可能的命令和参数
        detection_commands = [
            "iperf3",
            "iperf3.exe",
            "iperf",
            "iperf.exe"
        ]
        
        current_dir = os.path.dirname(os.path.abspath(__file__))
        
        # 首先检查用户选择的来源
        source = self.iperf_source_var.get()
        
        if source == "builtin":
            self.log("尝试使用内置版本的iperf3...")
            if self.extract_builtin_iperf3() and os.path.exists(self.builtin_iperf3_path):
                self.detected_iperf_command = self.builtin_iperf3_path
                self.log(f"成功使用内置iperf3: {self.builtin_iperf3_path}")
                return True
            else:
                self.log("无法使用内置版本，尝试使用系统安装的版本...")
        
        # 检查当前目录是否有可执行文件
        for cmd_name in detection_commands:
            full_path = os.path.join(current_dir, cmd_name)
            if os.path.exists(full_path) and os.access(full_path, os.X_OK):
                self.log(f"在当前目录找到可执行文件: {full_path}")
                self.detected_iperf_command = full_path
                return True
        
        # 然后尝试系统路径中的命令
        for cmd_name in detection_commands:
            try:
                self.log(f"尝试命令: {cmd_name} -v")
                # 使用shell=True在Windows上更可靠
                result = subprocess.run([cmd_name, "-v"], 
                                      stdout=subprocess.PIPE, 
                                      stderr=subprocess.PIPE, 
                                      timeout=3,
                                      universal_newlines=True,
                                      shell=True)
                
                self.log(f"命令返回码: {result.returncode}")
                stdout_lower = result.stdout.lower()
                stderr_lower = result.stderr.lower()
                
                # 检查是否真的执行了iperf命令
                if "'" + cmd_name + "' 不是内部或外部命令" in stderr_lower or \
                   "无法将\"" + cmd_name + "\"识别为" in stderr_lower:
                    self.log(f"命令不可用: {cmd_name}")
                    continue
                
                # 检查输出是否包含iperf相关信息
                if "iperf" in stdout_lower or "iperf" in stderr_lower:
                    self.log(f"检测成功: {cmd_name} 被识别为iperf")
                    self.detected_iperf_command = cmd_name
                    return True
                    
            except FileNotFoundError:
                self.log(f"命令未找到: {cmd_name}")
                continue
            except Exception as e:
                self.log(f"执行命令 {cmd_name} 时出错: {str(e)}")
                continue
        
        # 检查系统PATH
        path_env = os.environ.get('PATH', '')
        self.log(f"系统PATH: {path_env[:200]}...")
        
        self.log("所有检测方法均失败，iperf3可能未安装或未在系统PATH中")
        return False
        
    # log方法已在前面定义，删除重复的方法

if __name__ == "__main__":
    root = tk.Tk()
    app = Iperf3TesterApp(root)
    root.mainloop()